#include "StdAfx.h"
#include "PEImage.h"

CPEImage::CPEImage(const CAtlFile& file)
	: m_file(file)
	, m_base(NULL)
{
	if (SUCCEEDED(m_map.MapFile(m_file)))
	{
		m_base = static_cast<PBYTE>(m_map);
	}
}

CPEImage::CPEImage(const tstring& path)
	: m_base(NULL)
	, m_Managed(false)
	, m_Win32(false)
	, m_COM(false)
{
	if (SUCCEEDED(m_file.Create(path.c_str(),
								FILE_GENERIC_READ,
								FILE_SHARE_READ | FILE_SHARE_WRITE | FILE_SHARE_DELETE,
								OPEN_EXISTING)))
	{
		if (SUCCEEDED(m_map.MapFile(m_file)))
		{
			m_base = static_cast<PBYTE>(m_map);

			if (m_base != NULL)
			{
				Initialize();
			}
		}
	}
}

void CPEImage::Initialize()
{
	PIMAGE_DOS_HEADER pDosHeader = reinterpret_cast<PIMAGE_DOS_HEADER>(m_base);
	
	if (pDosHeader->e_magic == IMAGE_DOS_SIGNATURE && pDosHeader->e_lfanew != 0)
	{
		PIMAGE_NT_HEADERS32 pNtHeader32 = MakePtr<PIMAGE_NT_HEADERS32>(m_base, pDosHeader->e_lfanew);
		PIMAGE_NT_HEADERS64 pNtHeader64 = NULL;

		m_Win32 = (pNtHeader32->Signature == IMAGE_NT_SIGNATURE);

		if (m_Win32)
		{
			PIMAGE_COR20_HEADER pCorHeader;

			if (pNtHeader32->OptionalHeader.Magic == IMAGE_NT_OPTIONAL_HDR64_MAGIC)
			{
				pNtHeader64 = reinterpret_cast<PIMAGE_NT_HEADERS64>(pNtHeader32);
				pCorHeader = GetCor20Header(pNtHeader64);
				
				GetExports(pNtHeader64);
			}
			else
			{
				pCorHeader = GetCor20Header(pNtHeader32);
				
				GetExports(pNtHeader32);
			}

			bool reg = false;
			bool unreg = false;

			for (vector<CExportedFunction>::const_iterator it  = m_Exports.begin(); 
				                                           it != m_Exports.end();
														 ++it)
			{
				if (it->Name() == "DllRegisterServer")
				{
					reg = true;
					
					if (unreg) 
						break;
				}

				if (it->Name() == "DllUnregisterServer")
				{
					unreg = true;
					
					if (reg)
						break;
				}
			}
			
			m_COM = (reg && unreg);
			m_Managed = (pCorHeader != NULL);
		}
	}
}

template<class THeader>
void CPEImage::GetExports(const THeader* pNTHeader)
{
	const IMAGE_DATA_DIRECTORY* pDataDir = 
		GetDirectory(IMAGE_DIRECTORY_ENTRY_EXPORT, pNTHeader);
	
	if (pDataDir == NULL)
	{
		return;
	}

	PIMAGE_EXPORT_DIRECTORY pExportDir = 
		GetPtrFromRVA<IMAGE_EXPORT_DIRECTORY>(pDataDir->VirtualAddress, pNTHeader);

	if (pExportDir == NULL)
	{
		return;
	}

	PDWORD pdwFunctions = GetPtrFromRVA<DWORD>(pExportDir->AddressOfFunctions, pNTHeader);
	PDWORD pdwNames     = GetPtrFromRVA<DWORD>(pExportDir->AddressOfNames, pNTHeader);
	PWORD  pwOrdinals   = GetPtrFromRVA<WORD>(pExportDir->AddressOfNameOrdinals, pNTHeader);

	for (unsigned i = 0; i < pExportDir->NumberOfFunctions; ++i, ++pdwFunctions)
	{
		LPSTR szName = NULL;
		LPSTR szForwarder = NULL;
		DWORD rvaEntryPoint = *pdwFunctions;

		if (rvaEntryPoint == 0)
		{
			continue;
		}

		for (unsigned ord = 0; ord < pExportDir->NumberOfNames; ++ord)
		{
			if (pwOrdinals[ord] == i)
			{
				szName = GetPtrFromRVA<CHAR>(pdwNames[ord], pNTHeader);
				break;
			}
		}

		if ((rvaEntryPoint >= pDataDir->VirtualAddress) &&
			(rvaEntryPoint <= pDataDir->VirtualAddress + pDataDir->Size))
		{
			szForwarder = GetPtrFromRVA<CHAR>(rvaEntryPoint, pNTHeader);
		}

		m_Exports.push_back(CExportedFunction(rvaEntryPoint, i, szName, szForwarder));
	}
}

template<class THeader>
PIMAGE_COR20_HEADER CPEImage::GetCor20Header(const THeader* pNTHeader) const
{
	const IMAGE_DATA_DIRECTORY* directory 
		= GetDirectory(IMAGE_DIRECTORY_ENTRY_COM_DESCRIPTOR, pNTHeader);

	if (directory == NULL)
	{
		return NULL;
	}
	
	return GetPtrFromRVA<IMAGE_COR20_HEADER>(directory->VirtualAddress, pNTHeader);
}

template <class THeader> 
LPVOID CPEImage::GetPtrFromRVA(DWORD rva, const THeader* pNTHeader) const
{
	PIMAGE_SECTION_HEADER pSectionHdr = GetEnclosingSectionHeader(rva, pNTHeader);

	if (pSectionHdr == NULL)
	{
		return NULL;
	}

	DWORD delta = (pSectionHdr->VirtualAddress - pSectionHdr->PointerToRawData);
	return (m_base + rva - delta);
}

template <class THeader>
PIMAGE_SECTION_HEADER CPEImage::GetEnclosingSectionHeader(DWORD rva, const THeader* pNTHeader) const
{
    PIMAGE_SECTION_HEADER section = IMAGE_FIRST_SECTION(pNTHeader);
    
    for (unsigned i = 0; i < pNTHeader->FileHeader.NumberOfSections; i++, section++ )
    {
		DWORD size = section->Misc.VirtualSize;

		// Watcom's linker sets the Misc.VirtualSize field to 0
		if (size == 0)
		{
			size = section->SizeOfRawData;
		}

        // Is the RVA within this section?
        if ((rva >= section->VirtualAddress) && 
            (rva < (section->VirtualAddress + size)))
		{
            return section;
		}
    }
    
    return 0;
}